#!/bin/env python2
# -*- coding: utf-8 -*-
from __future__ import division
from collections import defaultdict
from itertools import product
from math import log
from utils import uopen, AttrDict

class AlignmentModel():
    def __init__(self, giza_model_path, params):
        """
        Example of usage:
        >>> model = AlignmentModel('/home/toto/my_giza_model')
        >>> my_alignment = model.align('Please align these sentences', 'Aligne ces phrases s\'il te plaît')
        >>> my_alignment.src_alignment
        [0, 0, 2, 3]

        The naming convention for the files of the model is:
            my_giza_model.{src}-{trg}.ttable
            my_giza_model.{src}-{trg}.atable
            my_giza_model.{src}.vcb
            
            my_giza_model.{trg}-{src}.ttable
            my_giza_model.{trg}-{src}.atable
            my_giza_model.{trg}.vcb

        See the documentation of the ``align`` method to understand how an alignment works.
        """
        params = AttrDict(**params)
        params.setdefault(giza_model_type='ibm2')
        
        src_lang, trg_lang = params.src_lang, params.trg_lang

        self.src_vcb = read_vcb('{}.{}.vcb'.format(giza_model_path, src_lang))
        self.trg_vcb = read_vcb('{}.{}.vcb'.format(giza_model_path, trg_lang))
        self.model_type = params.giza_model_type
      
        name = '{}.{{}}-{{}}.{{}}'.format(giza_model_path)
        self.src_trg_ttable = read_ttable(name.format(src_lang, trg_lang, 'ttable'))
        self.trg_src_ttable = read_ttable(name.format(trg_lang, src_lang, 'ttable'))

        if self.model_type == 'ibm2':
            self.src_trg_atable = read_atable(name.format(src_lang, trg_lang, 'atable'))
            self.trg_src_atable = read_atable(name.format(trg_lang, src_lang, 'atable'))
        else:
            self.src_trg_atable = None
            self.trg_src_atable = None

        self.src_trg_lexicon = build_lexicon(self.src_vcb, self.trg_vcb, self.src_trg_ttable, params)
        self.trg_src_lexicon = build_lexicon(self.trg_vcb, self.src_vcb, self.trg_src_ttable, params)

    @staticmethod
    def _align(src, trg, ttable, atable=None):
        src, trg = src[:], trg[:]
        lf, le = len(src), len(trg)
        # Insert NULL word
        src.insert(0, 0)

        # Translation table and alignment table
        if atable is not None:
            a = defaultdict(float, (((i, j), s) for i, j, s in atable[lf]))
        else:
            a = defaultdict(lambda: 1.0)

        t = dict()
        for n, m in product(src, trg):
            t[(n, m)] = ttable.get((n, m), 0)

        alignment = []
        log_proba = 0
        for j, m in enumerate(trg):
            # For each English word, take the French word with the highest translation probability
            # In the a table, the second column (trg) starts at 1. The 0 in the first column stands for NULL.
            p, i = max((a[(i, j + 1)] * t[(src[i], m)], i) for i in range(lf + 1))

            # Avoids null probabilities
            e = 1e-10 # arbitrary small number...
            if p < e:
                i = 0
                p = e

            log_proba -= log(p)
            alignment.append(i)
        return alignment, log_proba

    def align(self, src_string, trg_string):
        """
        Produce an IBM Model 2 alignment of source and target sentences.
        ``ttable`` and ``atable`` are respectively the translation and alignment tables,
        from source language to target language.
        
        Return the alignment as an array and its log probability.
        The alignment contains for each English word, the French word that is aligned to it.

        For example:
        >>> alignment = model.align('Please align these sentences', 'Aligne ces phrases s\'il te plaît')
        >>> alignment.src_alignment, alignment.src_score
        ([0, 0, 2, 3], 51.8437)

        Here the source language is English, and the target language French.

        Please [0] -> NULL [0]
        align [1] -> NULL [0]
        these [2] -> ces [2]
        sentences [3] -> phrases [3]

        The position 0 on the source side is reserved for the NULL word.
        """
        src_words, trg_words = src_string.split(), trg_string.split()
        src = [self.src_vcb.get(word) for word in src_words]
        trg = [self.trg_vcb.get(word) for word in trg_words]

        src_alignment, src_score = AlignmentModel._align(src, trg, self.src_trg_ttable, self.src_trg_atable)
        trg_alignment, trg_score = AlignmentModel._align(trg, src, self.trg_src_ttable, self.trg_src_atable)

        alignment = AttrDict(src_words=src_words, trg_words=trg_words,
                src_len=len(src), trg_len=len(trg),
                src_alignment=src_alignment,
                trg_alignment=trg_alignment,
                src_score=src_score,
                trg_score=trg_score)

        return alignment

def read_vcb(vcb_filename):
    """
    Read a vocabulary file in the Giza format and return a dictionary.
    """
    vcb = dict()
    with uopen(vcb_filename) as vcb_file:
        for line in vcb_file:
            index, word, _ = line.split()
            vcb[word] = int(index)

    return vcb

def read_ttable(ttable_filename):
    """
    Read the Giza translation table from disk.
    """
    with open(ttable_filename) as ttable_file:
        ttable = dict()
        for line in ttable_file:
            src_id, trg_id, score = line.split()
            ttable[(int(src_id), int(trg_id))] = float(score)

        return ttable

def read_atable(atable_filename): 
    """
    Read the Giza alignment table from disk.
    When building the alignment table, the column ``le`` (length of the English side) is ignored.
    """
    with open(atable_filename) as atable_file:
        # Bug with the 'le' column on mgiza (always 100)
        atable = defaultdict(list)
        for line in atable_file: # fr_pos en_pos lf le score
            values = line.split()
            fr_pos, en_pos, lf, le = map(int, values[:4])
            score = float(values[4])
            #atable[(lf, le)].append((fr_pos, en_pos, score))
            atable[lf].append((fr_pos, en_pos, score))

    return atable

def build_lexicon(src_vcb, trg_vcb, ttable, params):
    """
    Build a French-English lexicon with the Giza translation table.
    All word translations with a probability score above the threshold are kept.
    """
    params.setdefault(lexicon_threshold=0.1, lexicon_max=10)

    rev_trg_vcb = dict((v, k) for k, v in trg_vcb.iteritems())
    rev_src_vcb = dict((v, k) for k, v in src_vcb.iteritems())
   
    table = [(s, (n, m)) for (n, m), s in ttable.iteritems()
        if s >= params.lexicon_threshold and n != 0]

    table.sort(reverse=True)

    lexicon = defaultdict(list)
    for _, (n, m) in table:
        src_word, trg_word = rev_src_vcb[n], rev_trg_vcb[m]
        translations = lexicon[src_word]
        if len(translations) < params.lexicon_max:
            translations.append(trg_word)

    return lexicon

